# coding=utf-8
"""
需要：可执行算法(就是方法), 开放的mongodb地址
实现原理:
    一张发布任务表 ParentTask (每行对应一个总任务)
        {name: tasktable， status: going/end, log:[执行程序id]}
        {name: tasktable， status: going/end, log:[执行程序id]}
        {name: tasktable， status: going/end, log:[执行程序id]}
    

    若干张执行任务表，其名字在总任务表name字段里 tasktable

        {'_id': Objiectid(xxxx), status: new/going/end, object:要执行的算法, args: 执行的参数, result:执行的结果, sstart: 开始执行的时间，end: 执行完毕的时间}    （其中开始执行时间和结束执行时间时时间戳的形式，以后检测假死用）

        {'_id': Objiectid(xxxx), status: new/going/end, object:要执行的算法, args: 执行的参数, result:执行的结果, sstart: 开始执行的时间，end: 执行完毕的时间}    （其中开始执行时间和结束执行时间时时间戳的形式，以后检测假死用）

        {'_id': Objiectid(xxxx), status: new/going/end, object:要执行的算法, args: 执行的参数, result:执行的结果, sstart: 开始执行的时间，end: 执行完毕的时间}    （其中开始执行时间和结束执行时间时时间戳的形式，以后检测假死用）

    
有了以上两类表就可以执行任务了，流程如下:
1.去总任务表看看有没有going状态任务，如果有再看看执行日志里有没有自己的id如果没有则获取任务表名字开始执行
2.拿到任务表名之后会对应已经选好的参数表，原子检测是否有new状态任务有的话将其状态改变为going开始执行
3.执行结束返回结果并将状态改为end
4.继续寻找可执行任务，如果发现没有new任务则全局检测一下任务是否都是end，如果都是end就将总任务状态改成end
5.循环 1.

"""
import gc
import json
from socket import *
import random
import os
import time
import pickle
import signal
from pymongo import MongoClient
from bson.objectid import ObjectId


def get_host_ip():
    try:
        s = socket(AF_INET, SOCK_DGRAM)
        s.connect(('8.8.8.8', 80))
        ip = s.getsockname()[0]
    finally:
        s.close()

    return ip

class GetTask(object):
    """获取任务类"""

    def __init__(self, addr, pid, db, table):
        self.conn = MongoClient( addr[0], addr[1])
        self.db = self.conn[db][table]
        self.pid = pid


    def main(self):
        """循环检测任务队列"""

        while True:
            tasks = self.db.find()
            for task in tasks:
                if task['status'] == 'going':
                    if self.pid not in task['log']:
                        print '查询到一个尚未执行过的任务,进行注册返回程序'
                        name = task['name']
                        ids = task['_id']
                        self.db.update({'_id': ids}, {'$push': {'log':{'pid': self.pid, 'start':time.time(), 'ip':get_host_ip()}}})
                        self.conn.close()
                        return name, ids
            time.sleep(1)


class RunTask(object):
    """得到任务列表开始执行任务"""
    def __init__(self, addr, pid, db, table, worker_id):
        self.conn = MongoClient( addr[0], addr[1])
        self.db = self.conn[db][table]
        self.pid = pid
        self.worker = self.conn[db]['worker']
        self.worker_id = worker_id
        

    def get_one_task(self):
        """得到一个任务对象"""
        while True:
            task = self.db.find_one({"status": 'new'})
            # 如果不存在这个条件的对象，说明所有任务执行完毕
            if not task:
                return task
            # 存在这个对象，则开始修改任务状态
            else:
                task = self.db.find_and_modify({'_id':task['_id']}, {'$set':{'status':'going', 'log':self.pid, 'start':time.time(), 'ip': get_host_ip()}})
                # 如果有其他任务提前就改了状态 repetion
                if task:
                    return task
                

    def run(self, task):
        """run process"""
        try:
            colony_fun = pickle.loads(task['def'])
            task['result'] = colony_fun(task['args'])
        except Exception as e:
            print 'run task error:', e
            task['result'] = e
            task['status'] = 'error'
        else:
            task['status'] = 'end'
        finally:
            info = self.db.find_and_modify({'_id': task['_id']},{'$set': {'result': task['result'], 'status': task['status'], 'end': time.time()}})
            info = self.db.find_one({"_id": info["_id"]})
            info['task_parent'] = str(self.db)
            print info
            self.worker.update({'_id': self.worker_id}, {'$push':{'log':info}})
        #print 'Mission accomplished   id:',task['_id'], 'details:', str(self.db)

    def main(self):
        """循环执行"""
        while True:
            task = self.get_one_task()
            if not task:
                print 'not have new task    Finish'
                self.conn.close()
                return
            self.run(task)

        

def checkandupdate(addr, myselfpid, db, table, task_id, parent):
        conn = MongoClient( addr[0], addr[1])
        parent = conn[db][parent]
        db = conn[db][table]
        
        task = db.find_one({'status': 'going'})
        
        if not task:
            parent.find_and_modify({'_id': task_id}, {"$set":{'status': 'end', 'end': time.time()}})
            print 'task name :',table, 'run end'
 

def heartbeat(db, addr, worker_id):

    tcp_client_socket = socket(AF_INET, SOCK_STREAM)

    server_ip = "0.0.0.0"
    server_port = 25687
    tcp_client_socket.connect((server_ip, server_port))
    data = json.dumps({'db': db, 'id': str(worker_id), 'addr': addr})
    
    send_data = data

    tcp_client_socket.send(send_data.encode("utf-8"))
    print '发送注册信息'
    return tcp_client_socket

def worker_main():
    """run process"""
    
    myselfpid = os.getpid() + round(random.random(),3)  # 尽可能不和其他电脑重叠
    
    database_db = 'Colony'  #  使用的数据库名字
    parent_table = 'ParentTable'  #　使用总任务表名字
    addr = ('0.0.0.0', 27017)  # mongdb地址
    
    conn = MongoClient(addr[0], addr[1])
    db = conn[database_db]['worker']
    worker_id = db.insert({'pid': int(myselfpid), 'ip': get_host_ip(), 'log': list(), 'status': 'going', 'start':str(time.strftime("%Y-%m-%d %H:%M %p", time.localtime()))})
    
    # 执行心跳注册
    hb = heartbeat(database_db, addr, worker_id)
    conn.close()
    try:
        while True:
            # 获取任务和表id
            gettask = GetTask(addr, myselfpid, database_db, parent_table)
            task, task_id = gettask.main()
            
            # 执行任务
            runtask = RunTask(addr, myselfpid, database_db, task, worker_id)
            runtask.main()

            # 检测parent任务是否全部执行完毕，如果你是最后一个执行完的就把总任务状态改成 end
            checkandupdate(addr, myselfpid, database_db, task, task_id, parent_table)

            time.sleep(1)
    except Exception as e:
        print e
 #   db.update({'_id': worker_id}, {'$set':{'status': 'end', 'end': str(time.strftime("%Y-%m-%d %H:%M %p", time.localtime()))}})
#-------此处放置自定义算法需要的模块和算法----------
# 因为时第一版比较简陋 注意算法和脚本本来的方法不要重名


#---------------------------------------------------









if __name__ == "__main__":
    worker_main()

